"""
Callbacks for integrating continual learning with HuggingFace models.
"""

import os
from typing import Optional

from ..task_callbacks import TaskCallback


class HuggingFaceModelCallback(TaskCallback):
    """
    Callback for integration with Hugging Face Hub.

    This callback enables saving models to Hugging Face Hub at various points
    during training, such as after each task, epoch, or at the end of training.

    Args:
        save_to_hub (bool): Whether to save to Hugging Face Hub
        hub_repo_id (str): Repository ID on Hugging Face Hub
        hub_token (str): API token for Hugging Face Hub
        save_every_n_epochs (int): Save to Hub every N epochs (0 to disable)
        push_to_hub_on_task_end (bool): Save to Hub when a task ends
    """

    def __init__(
        self,
        save_to_hub: bool = False,
        hub_repo_id: Optional[str] = None,
        hub_token: Optional[str] = None,
        save_every_n_epochs: int = 0,
        push_to_hub_on_task_end: bool = True,
    ):
        super().__init__()
        self.save_to_hub = save_to_hub
        self.hub_repo_id = hub_repo_id
        self.hub_token = hub_token
        self.save_every_n_epochs = save_every_n_epochs
        self.push_to_hub_on_task_end = push_to_hub_on_task_end

        # Check if we can save to hub
        try:
            from huggingface_hub import HfApi

            self.hf_api = HfApi()
            self.can_save_to_hub = save_to_hub and hub_repo_id is not None
        except ImportError:
            print("Warning: huggingface_hub package not found. Cannot save to Hub.")
            self.can_save_to_hub = False

    def on_train_epoch_end(self, trainer, pl_module):
        """Save model to Hub at epoch end if configured."""
        if not self.can_save_to_hub or self.save_every_n_epochs <= 0:
            return

        current_epoch = trainer.current_epoch

        if current_epoch % self.save_every_n_epochs == 0:
            self._save_model_to_hub(trainer, pl_module, f"epoch_{current_epoch}")

    def on_task_end(self, trainer, pl_module, task_id, dataloader=None):
        """Save model to Hub when a task ends."""
        super().on_task_end(trainer, pl_module, task_id, dataloader)

        if not self.can_save_to_hub or not self.push_to_hub_on_task_end:
            return

        self._save_model_to_hub(trainer, pl_module, f"task_{task_id}_complete")

    def on_train_end(self, trainer, pl_module):
        """Save final model to Hub."""
        if not self.can_save_to_hub:
            return

        self._save_model_to_hub(trainer, pl_module, "final_model")

    def _save_model_to_hub(self, trainer, pl_module, tag):
        """Save the model to Hugging Face Hub."""
        try:
            from huggingface_hub import create_repo
            from transformers import PreTrainedModel
            import tempfile

            # Check if repository exists, create if not
            try:
                self.hf_api.repo_info(repo_id=self.hub_repo_id, token=self.hub_token)
            except:
                create_repo(
                    repo_id=self.hub_repo_id,
                    token=self.hub_token,
                    private=False,
                    exist_ok=True,
                )

            # Create temporary directory
            with tempfile.TemporaryDirectory() as tmpdirname:
                # Save model
                if hasattr(pl_module.model, "backbone") and hasattr(
                    pl_module.model.backbone, "model"
                ):
                    if isinstance(pl_module.model.backbone.model, PreTrainedModel):
                        # Save HuggingFace model
                        pl_module.model.backbone.model.save_pretrained(tmpdirname)

                        # Save model card
                        with open(os.path.join(tmpdirname, "README.md"), "w") as f:
                            f.write(f"# {self.hub_repo_id}\n")
                            f.write(
                                f"This model was trained using continual learning.\n"
                            )
                            f.write(f"Tag: {tag}\n")
                            f.write(f"Current task: {pl_module.current_task}\n")

                # Upload to Hub
                self.hf_api.upload_folder(
                    folder_path=tmpdirname,
                    repo_id=self.hub_repo_id,
                    token=self.hub_token,
                    commit_message=f"Upload model - {tag}",
                )

                print(
                    f"Model saved to HuggingFace Hub: {self.hub_repo_id} (tag: {tag})"
                )
        except Exception as e:
            print(f"Error saving model to Hub: {e}")
